diff --git a/mesh_renderer/rasterize_triangles.py b/mesh_renderer/rasterize_triangles.py
index 813c864..8b6d720 100644
--- a/mesh_renderer/rasterize_triangles.py
+++ b/mesh_renderer/rasterize_triangles.py
@@ -104,7 +104,7 @@ def rasterize_clip_space(clip_space_vertices, attributes, triangles,
   if len(clip_space_vertices.shape) != 3:
     raise ValueError('The vertex buffer must be 3D.')
 
-  vertex_count = clip_space_vertices.shape[1].value
+  vertex_count = clip_space_vertices.shape[1]
 
   batch_size = tf.shape(clip_space_vertices)[0]
   
@@ -126,7 +126,7 @@ def rasterize_clip_space(clip_space_vertices, attributes, triangles,
         b, tf.reshape(barycentric_coords, [-1, 3]))
 
     vertex_ids = tf.gather(triangles, tf.reshape(triangle_ids, [-1]))
-    reindexed_ids = tf.add(vertex_ids, b * clip_space_vertices.shape[1].value)
+    reindexed_ids = tf.add(vertex_ids, b * clip_space_vertices.shape[1])
     per_image_vertex_ids = per_image_vertex_ids.write(b, reindexed_ids)
 
     return b+1, per_image_barycentric_coordinates, per_image_vertex_ids
